package com.riche.aicodehelper.ai.guardrail;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;
import org.springframework.context.annotation.Configuration;

import java.util.Set;

/**
 * @ClassName: SafeInputGuardrail
 * @Description: 护轨输入类
 * @Author: Gaoruiqi
 * @Date: 2025-09-27 20:45
 * @Version: 1.0
 **/

public class SafeInputGuardrail implements InputGuardrail {

    // 敏感词集合
    private static final Set<String> sensitiveWords = Set.of("kill", "evil");

    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        // 获取用户输入并转换为小写确保不包含敏感词汇
        String inputText = userMessage.singleText().toLowerCase();
        String[] words = inputText.split("\\W+");
        for (String word : words) {
            if (sensitiveWords.contains(word)) {
                return fatal("输入包含敏感词汇，请勿输入敏感词汇");
            }
        }
        return success();
    }
}
